package com.xjai.utils;

import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * Server Send Event SSE工具类。
 */

@Slf4j
public class SSEServer {

    /**
     * 使用map对象，关联ID 和SSE的服务链接
     */
    private static Map<String, SseEmitter> sseClients = new ConcurrentHashMap<>();

    public static SseEmitter connect(String userId){
        // 设置超市时间，0 代表永不过期， 默认30秒，超时未完成任务则会抛出异常
        SseEmitter sseEmitter = new SseEmitter(0L);

        // 注册SSE的回调方法
        sseEmitter.onCompletion(completionCallBack(userId));
        sseEmitter.onError(errorCallBack(userId));
        sseEmitter.onTimeout(timeoutCallBack(userId));
        sseClients.put(userId, sseEmitter);
        return  sseEmitter;

    }

    private static Runnable timeoutCallBack(String userId) {
        return () -> {
            log.info("SSE链接超时，用户ID为: {}", userId);
            removeConnection(userId);
        };
    }

    private static Consumer<Throwable> errorCallBack(String userId) {
        return Throwable -> {
          log.info("SSE链接发生错误，用户ID为: {}", userId);
          removeConnection(userId);
        };
    }

    /**
     * SSE链接完成后的回调方法（关闭链接的时候调用）
     * @param userId
     * @return
     */
    private static Runnable completionCallBack(String userId) {
        return () -> {
            log.info( "SSE链接完成并结束， 用户ID为: {}", userId );
            removeConnection(userId);
        };
    }

    private static void removeConnection(String userId) {
        sseClients.remove(userId);
        log.info( "SSE链接被移除， 移除的用户ID为: {}", userId );
    }

    public static void sendMessage(String userId,String message, String type){
        if(CollectionUtils.isEmpty(sseClients)){
            return;
        }

        if(sseClients.containsKey(userId)){
            SseEmitter sseEmitter = sseClients.get(userId);
            sendEmitterMessage(sseEmitter, userId, message, type);
        }

    }

    private static void sendEmitterMessage(SseEmitter sseEmitter, String userId, String message, String add) {
        SseEmitter.SseEventBuilder data = SseEmitter.event().id(userId).name(add).data(message);
        try {
            sseEmitter.send(data);
        } catch (IOException e) {
            log.error( "用户[{}]的消息推送发生异常", userId );
            removeConnection(userId);
        }
    }


}
